# !/usr/bin/env python3
# -*- coding: utf-8 -*-
# @Time : 2023/4/23 22:15
# @Author : xyjiang1216
# @Email : jiangxuyan@nudt.edu.cn
# @File : construct_ring_topo_graph.py
# @Version: OpenPlanner2.0
from random import choice

import networkx as nx

from OpenPlanner.core.src.topo_construction_strategy.topo_interface import TopoStrategy
from OpenPlanner.util.topo_graph_utils import show_topology_graph


class RingTopoStrategy(TopoStrategy):
    # 重写方法
    def generate_topo_graph(self):
        # print('环形拓扑策略')
        # print(self.topo_generation_param.topo_type)

        sw_num = self.topo_param.sw_num
        es_num_per_sw_set = self.topo_param.es_num_per_sw_set

        G = nx.DiGraph()

        # 第一步，添加所有的交换机节点
        sw_id = 0
        for sw in range(sw_num):
            G.add_node(sw_id, node_id=sw_id, node_type='SW', es_set=[], current_avail_port=0)
            sw_id += 1

        # 第二步，添加所有的终端节点
        es_id = sw_num
        for sw in range(sw_num):
            try:
                es_num_of_current_sw = choice(es_num_per_sw_set)
            except IndexError:
                print('配置文件中的es_num_per_sw_set不能为空！')
                exit(0)
            for es in range(es_num_of_current_sw):
                G.add_node(es_id, node_id=es_id, node_type='ES', current_avail_port=0)
                # 这个节点应当连接到哪一个交换机
                es_set = G.nodes[sw]['es_set']
                es_set.append(es_id)
                G.nodes[sw]['es_set'] = es_set
                es_id += 1

        # 第三步，添加交换机节点之间的链路
        sw_id_set = range(0, sw_num)
        link_id = 0
        for sw in sw_id_set:
            src_port = G.nodes[sw]['current_avail_port']
            dst_port = G.nodes[(sw + 1) % sw_num]['current_avail_port']
            G.add_edge(sw, (sw + 1) % sw_num, link_id=link_id, src_port=src_port, dst_port=dst_port)
            link_id += 1
            src_port = G.nodes[(sw + 1) % sw_num]['current_avail_port']
            dst_port = G.nodes[sw]['current_avail_port']
            G.add_edge((sw + 1) % sw_num, sw, link_id=link_id, src_port=src_port, dst_port=dst_port)
            link_id += 1
            G.nodes[sw]['current_avail_port'] += 1
            G.nodes[(sw + 1) % sw_num]['current_avail_port'] += 1

        # 第四步，添加终端节点与交换机节点之间的链路
        for sw in sw_id_set:
            es_set = G.nodes[sw]['es_set']
            for es in es_set:
                src_port = G.nodes[sw]['current_avail_port']
                dst_port = G.nodes[es]['current_avail_port']
                G.add_edge(sw, es, link_id=link_id, src_port=src_port, dst_port=dst_port)
                link_id += 1
                src_port = G.nodes[es]['current_avail_port']
                dst_port = G.nodes[sw]['current_avail_port']
                G.add_edge(es, sw, link_id=link_id, src_port=src_port, dst_port=dst_port)
                link_id += 1
                G.nodes[es]['current_avail_port'] += 1
                G.nodes[sw]['current_avail_port'] += 1

        # 如果要显示拓扑
        show_topology_graph(G)

        # print(G)
        # 返回生成的拓扑图
        return G